﻿// dllmain.cpp : 定义 DLL 应用程序的入口点。
#include "pch.h"
#include <windows.h>
#include <detours.h>
#include <fstream>
#include <mutex>
#include <string>
#include <sstream>
#include <locale>
#include <codecvt>
#include <iostream>
#include <algorithm>      // C++11
#include <unordered_set>  // C++11
#include "dllmain.h"

std::wstring ExtractProcessName(const std::wstring& fullCmd) {
    std::wstring exePath;

    // 处理带引号的路径（如 "C:\Program Files\cl.exe"）
    if (!fullCmd.empty() && fullCmd[0] == L'"') {
        size_t closingQuote = fullCmd.find(L'"', 1);
        if (closingQuote != std::wstring::npos) {
            exePath = fullCmd.substr(1, closingQuote - 1);
        }
    }
    else {
        size_t firstSpace = fullCmd.find(L' ');
        exePath = (firstSpace != std::wstring::npos) ?
            fullCmd.substr(0, firstSpace) : fullCmd;
    }

    // 提取文件名部分（如 cl.exe）
    size_t lastSlash = exePath.find_last_of(L"\\/");
    std::wstring exeName = (lastSlash != std::wstring::npos) ?
        exePath.substr(lastSlash + 1) : exePath;

    // 转换为小写（C++11 lambda）
    std::transform(exeName.begin(), exeName.end(), exeName.begin(),
        [](wchar_t c) { return towlower(c); });  

    return exeName;
}

    // 新增过滤逻辑
static const std::unordered_set<std::wstring> allowedNames = {
    L"cl.exe",  // Visual C++ 编译器
    L"g++.exe",     // GNU C++ 编译器
    L"gcc.exe",     // GNU C 编译器
    L"clang.exe",   // Clang 编译器
    L"clang++.exe", // Clang++ 编译器
    L"arm-none-eabi-gcc.exe",
    L"cc1.exe",
    L"link.exe" // 链接器
};

// 使用 WinAPI 判断文件是否存在（C++14 兼容）
bool FileExists(const std::wstring& filePath) {
    DWORD attr = GetFileAttributesW(filePath.c_str());
    return (attr != INVALID_FILE_ATTRIBUTES && !(attr & FILE_ATTRIBUTE_DIRECTORY));
}

// 原始函数指针
static BOOL(WINAPI* TrueCreateProcessW)(
    LPCWSTR lpApplicationName,
    LPWSTR lpCommandLine,
    LPSECURITY_ATTRIBUTES lpProcessAttributes,
    LPSECURITY_ATTRIBUTES lpThreadAttributes,
    BOOL bInheritHandles,
    DWORD dwCreationFlags,
    LPVOID lpEnvironment,
    LPCWSTR lpCurrentDirectory,
    LPSTARTUPINFOW lpStartupInfo,
    LPPROCESS_INFORMATION lpProcessInformation
    ) = CreateProcessW;

// 原始CreateProcessA函数指针
static BOOL(WINAPI* TrueCreateProcessA)(
    LPCSTR lpApplicationName,
    LPSTR lpCommandLine,
    LPSECURITY_ATTRIBUTES lpProcessAttributes,
    LPSECURITY_ATTRIBUTES lpThreadAttributes,
    BOOL bInheritHandles,
    DWORD dwCreationFlags,
    LPVOID lpEnvironment,
    LPCSTR lpCurrentDirectory,
    LPSTARTUPINFOA lpStartupInfo,
    LPPROCESS_INFORMATION lpProcessInformation
    ) = CreateProcessA;

void WriteLog(const std::wstring& cmd) {
    // 创建或打开命名互斥体（确保全局唯一性）
    HANDLE hMutex = CreateMutexW(NULL, FALSE, L"Global\\IKGET_LogMutex");
    if (hMutex == NULL) return;
    // 等待获取互斥体所有权
    DWORD waitResult = WaitForSingleObject(hMutex, INFINITE);  // 5秒超时
    if (waitResult != WAIT_OBJECT_0) {
        // 处理超时或错误
        CloseHandle(hMutex);
        return;
    }
	//wprintf(L"[LOG] %s\n", cmd.c_str());
	wchar_t dir[MAX_PATH];
    GetEnvironmentVariableW(L"IKGET_WORKING_DIR", dir, MAX_PATH);
    std::wstring logPath = dir;
    logPath += L"\\compile_commands.txt";
    //wprintf(L"dll logPath: %s\n", logPath.c_str());

    bool fileExists = FileExists(logPath);

    // 使用二进制模式打开，避免 std::wofstream 插手编码
    std::wofstream logFile(logPath, std::ios::app | std::ios::binary);
    if (!logFile.is_open()) return;

    // 设置 UTF-16 Little Endian 编码
    logFile.imbue(std::locale(logFile.getloc(),
        new std::codecvt_utf16<wchar_t, 0x10ffff, std::little_endian>));

    // 如果文件不存在（即刚创建），写入 UTF-16 BOM（0xFEFF）
    if (!fileExists) {
        logFile.put(0xFEFF);  // BOM
    }
    logFile << GetCurrentProcessId() << L";" << cmd << std::endl;
    // 释放互斥体
    ReleaseMutex(hMutex);
    CloseHandle(hMutex);
}

// 记录命令行
void WriteCompileOptionsToLog(std::wstring& fullCmd)
{
    if (!fullCmd.empty()) {
        std::wstring lowercase_fullCmd = fullCmd;
		std::transform(lowercase_fullCmd.begin(), lowercase_fullCmd.end(), lowercase_fullCmd.begin(),
			[](wchar_t c) { return towlower(c); });
        //wprintf(L"[LOG fullCmd2] %s\n", fullCmd.c_str());
        std::wstring processName = ExtractProcessName(fullCmd);
        //wprintf(L"[LOG processName] %s\n", processName.c_str());
        if (allowedNames.find(processName) != allowedNames.end()) {
            //WriteLog(fullCmd);
            size_t atPos = lowercase_fullCmd.find(L"cl.exe\" @");
            if (atPos != std::wstring::npos) {
                std::wstring compilerPath = fullCmd.substr(0, atPos + wcslen(L"cl.exe\""));
                size_t filePathStart = atPos + wcslen(L"cl.exe\" @");
                std::wstring rspFilePath = fullCmd.substr(filePathStart);
                std::wifstream rspFile(rspFilePath, std::ios::in | std::ios::binary);
                if (rspFile.is_open()) {
                    rspFile.imbue(std::locale(rspFile.getloc(),
                        new std::codecvt_utf16<wchar_t, 0x10ffff, std::little_endian>));

                    // 检查 BOM（字节顺序标记）
                    wchar_t bom;
                    rspFile.read(&bom, 1);
                    if (bom == 0xFEFF) {
                        // 处理 UTF-16 编码文件，跳过 BOM
                        rspFile.seekg(2, std::ios::beg);
                    }
                    else {
                        // 如果没有 BOM，文件被视为 UTF-8 或其他编码
                        rspFile.seekg(0, std::ios::beg);
                    }

                    std::wstringstream fileContent;
                    fileContent << rspFile.rdbuf();
                    std::wstring fileText = fileContent.str();

                    WriteLog(compilerPath + L";" + fileText);
                }
                else {
                    WriteLog(L"Failed to open rsp file: " + rspFilePath);
                }
			}
            else {
                // 处理没有 @ 的情况
                WriteLog(fullCmd);
            }
        }
    }
}

// Hook 后的 CreateProcessW
BOOL WINAPI HookedCreateProcessW(
    LPCWSTR lpApplicationName,
    LPWSTR lpCommandLine,
    LPSECURITY_ATTRIBUTES lpProcessAttributes,
    LPSECURITY_ATTRIBUTES lpThreadAttributes,
    BOOL bInheritHandles,
    DWORD dwCreationFlags,
    LPVOID lpEnvironment,
    LPCWSTR lpCurrentDirectory,
    LPSTARTUPINFOW lpStartupInfo,
    LPPROCESS_INFORMATION lpProcessInformation) 
{
    BOOL result = FALSE;
    std::wstring fullCmd;
    //if (lpApplicationName) {
    //    fullCmd += L"\"";
    //    fullCmd += lpApplicationName;
    //    fullCmd += L"\"";
    //}
    if (lpCommandLine) {
        fullCmd = lpCommandLine;
    }
    //wprintf(L"[LOG fullCmd1] %s\n", fullCmd.c_str());

    wchar_t injectedEnv[MAX_PATH];

    // 获取 DLL 路径并注入到新的进程中
    wchar_t dllDir[MAX_PATH];
   //GetModuleFileNameW(NULL, dllPath, MAX_PATH);
    GetEnvironmentVariableW(L"IKGET_DLL_DIR", dllDir, MAX_PATH);
    std::wstring dllPath = dllDir;
    dllPath += L"\\ikgethook.dll";
    //wprintf(L"dll dllPath: %s\n", dllPath.c_str());

    //wcscpy_s(dllPath, MAX_PATH, L"e:\\work\\ikget_win\\ikgethook.dll");
    char dllPathA[MAX_PATH];
    //wprintf(L"[DLL PATH] %s\n", dllPath.c_str());
    WideCharToMultiByte(CP_ACP, 0, dllPath.c_str(), -1, dllPathA, MAX_PATH, NULL, NULL);

    if (GetEnvironmentVariableW(L"IKGET_INJECTED", injectedEnv, MAX_PATH) > 0) {
        wprintf(L"[TrueCreateProcessW] %s\n", fullCmd.c_str());
        //如果已注入，则跳过注入操作，防止递归
        WriteCompileOptionsToLog(fullCmd);
        result = TrueCreateProcessW(
            lpApplicationName,
            lpCommandLine,
            lpProcessAttributes,
            lpThreadAttributes,
            bInheritHandles,
            dwCreationFlags,
            lpEnvironment,
            lpCurrentDirectory,
            lpStartupInfo,
            lpProcessInformation
        );
        SetEnvironmentVariableW(L"IKGET_INJECTED", NULL);
        return result;
    }

    // 设置环境变量标志，防止递归
    SetEnvironmentVariableW(L"IKGET_INJECTED", L"1");
    
    result = DetourCreateProcessWithDllExW(
        lpApplicationName,
        lpCommandLine,
        lpProcessAttributes,
        lpThreadAttributes,
        bInheritHandles,
        dwCreationFlags,
        lpEnvironment,
        lpCurrentDirectory,
        lpStartupInfo,
        lpProcessInformation,
        dllPathA,
        NULL
    );
    // 设置环境变量标志，防止递归
    SetEnvironmentVariableW(L"IKGET_INJECTED", NULL);
    return result;
}

// Hook后的CreateProcessA
BOOL WINAPI HookedCreateProcessA(
    LPCSTR lpApplicationName,
    LPSTR lpCommandLine,
    LPSECURITY_ATTRIBUTES lpProcessAttributes,
    LPSECURITY_ATTRIBUTES lpThreadAttributes,
    BOOL bInheritHandles,
    DWORD dwCreationFlags,
    LPVOID lpEnvironment,
    LPCSTR lpCurrentDirectory,
    LPSTARTUPINFOA lpStartupInfo,
    LPPROCESS_INFORMATION lpProcessInformation)
{
    BOOL result = FALSE;
    // 将ANSI参数转换为宽字符
    std::wstring wAppName, wCmdLine;
    if (lpApplicationName) {
        int size = MultiByteToWideChar(CP_ACP, 0, lpApplicationName, -1, NULL, 0);
        wAppName.resize(size);
        MultiByteToWideChar(CP_ACP, 0, lpApplicationName, -1, &wAppName[0], size);
    }
    if (lpCommandLine) {
        int size = MultiByteToWideChar(CP_ACP, 0, lpCommandLine, -1, NULL, 0);
        wCmdLine.resize(size);
        MultiByteToWideChar(CP_ACP, 0, lpCommandLine, -1, &wCmdLine[0], size);
        wCmdLine.pop_back();
    }
    std::wstring fullCmd;
;
    if (!wCmdLine.empty()) {
        fullCmd = wCmdLine;
    }
    //wprintf(L"[LOG fullCmd A1] %s\n", fullCmd.c_str());

    wchar_t injectedEnv[MAX_PATH];

    // 获取 DLL 路径并注入到新的进程中
    wchar_t dllDir[MAX_PATH];
    //GetModuleFileNameW(NULL, dllPath, MAX_PATH);
    GetEnvironmentVariableW(L"IKGET_DLL_DIR", dllDir, MAX_PATH);
    std::wstring dllPath = dllDir;
    dllPath += L"\\ikgethook.dll";
    //wprintf(L"dll dllPath: %s\n", dllPath.c_str());

    //wcscpy_s(dllPath, MAX_PATH, L"e:\\work\\ikget_win\\ikgethook.dll");
    char dllPathA[MAX_PATH];
    //wprintf(L"[DLL PATH] %s\n", dllPath.c_str());
    WideCharToMultiByte(CP_ACP, 0, dllPath.c_str(), -1, dllPathA, MAX_PATH, NULL, NULL);

    if (GetEnvironmentVariableW(L"IKGET_INJECTED", injectedEnv, MAX_PATH) > 0) {
        wprintf(L"[TrueCreateProcessA] %s\n", fullCmd.c_str());
        //如果已注入，则跳过注入操作，防止递归
        WriteCompileOptionsToLog(fullCmd);
        return TrueCreateProcessA(
            lpApplicationName,
            lpCommandLine,
            lpProcessAttributes,
            lpThreadAttributes,
            bInheritHandles,
            dwCreationFlags,
            lpEnvironment,
            lpCurrentDirectory,
            lpStartupInfo,
            lpProcessInformation
        );
    }
    // 设置环境变量标志，防止递归
    SetEnvironmentVariableW(L"IKGET_INJECTED", L"1");

    result = DetourCreateProcessWithDllExA(
        lpApplicationName,
        lpCommandLine,
        lpProcessAttributes,
        lpThreadAttributes,
        bInheritHandles,
        dwCreationFlags,
        lpEnvironment,
        lpCurrentDirectory,
        lpStartupInfo,
        lpProcessInformation,
        dllPathA,  // 已转换为ANSI路径
        NULL
    );
    // 设置环境变量标志，防止递归
    SetEnvironmentVariableW(L"IKGET_INJECTED", NULL);
    return result;
}

// DLL 入口点
BOOL APIENTRY DllMain(HMODULE hModule, DWORD  reason, LPVOID reserved) {
    SetEnvironmentVariableW(L"IKGET_INJECTED", NULL);
    switch (reason) {
    case DLL_PROCESS_ATTACH:
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourAttach(&(PVOID&)TrueCreateProcessW, HookedCreateProcessW);
        DetourAttach(&(PVOID&)TrueCreateProcessA, HookedCreateProcessA);
        DetourTransactionCommit();
        break;
    case DLL_PROCESS_DETACH:
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourDetach(&(PVOID&)TrueCreateProcessW, HookedCreateProcessW);
        DetourDetach(&(PVOID&)TrueCreateProcessA, HookedCreateProcessA);
        DetourTransactionCommit();
        break;
    case DLL_THREAD_ATTACH:
        break;
    case DLL_THREAD_DETACH:
        break;
    }
    return TRUE;
}
